#include "master.h"
#include "mpi_partitionings.h"
#include "task_list.h"
#include "defines.h"

#include "raytrace.h"


#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/stat.h>

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void masterMain( ConfigData *data ) {
  char  *render_name;    // Filename where we save the rendered image.
  float *samples;       // Where we store the ray trace samples.
  float *recv_arr;
  int   *recv_arr_int;

  // Get information from the ConfigData structure passed in
  int rank            = data->mpi_rank;
  int num_procs       = data->mpi_procs;          // Used for number of strips/blocks
  PartType part_mode  = data->part_mode;
  
  int block_width     = data->block_width;        // Used for dynamic allocation.
  int block_height    = data->block_height;       // Used for dynamic allocation.
  int strip_height    = data->strip_height;       // Used for cyclical allocation.

  unsigned int height = data->height;    // How many pixels high to make the rendered image.
  unsigned int width  = data->width;     // How many pixels wide to make the rendered image.
  
  unsigned int part_size = ceil(sizeof(float) * width * height * 3);  
  
  unsigned int x, y;
  
  int node, count, pixel;
  
  unsigned int recv_addr;
  unsigned int recv_data_len;

  MPI_Status status;

  double comp_start, comp_stop, comp_time, tmp_comp_time;
  double comm_start, comm_stop, comm_time;
  int longest_running_node = 0;
  double total_comp_time;
  
  // Control variables
  unsigned int start_height, end_height;
  unsigned int start_width, end_width;
  
  
  // Task list
  task_list* my_tasks;
  int is_first_task = 1;
  int number_tasks = 0;  
  task_list* iter;
 
  // Give some information to the log

  printf ("(%d) Starting!\n", rank);

  // Verify that the 'renders' directory exists. Create it if it doesn't.
  struct stat stat_buf;
  stat( "renders", &stat_buf );
  if (!S_ISDIR(stat_buf.st_mode)) {
    if (mkdir( "renders", 0777 ) != 0) {
      fprintf( stderr, "Could not create 'renders' directory!\n" );
      fprintf( stderr, "Don't know where to save rendered images!\n" );
      exit(1);
    }
  }

  //////////////////////////////////////////////////////////////////////////////
  // Ray trace every point.
  //////////////////////////////////////////////////////////////////////////////

  switch (part_mode) {
    case PART_STATIC_STRIPS:
    
      samples = (float*) malloc( part_size );

      comp_start = MPI_Wtime();
      partition_static_strips_func (samples, data);
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;
      
      break;
      
    case PART_STATIC_BLOCKS:
      samples = (float*) malloc( part_size );

      comp_start = MPI_Wtime();
      partition_static_blocks_func (samples, data, &my_tasks);
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;
      
      break;

    case PART_STATIC_CYCLES:

      samples = (float*) malloc( part_size );
     
      comp_start = MPI_Wtime();
      partition_static_cycles_func (samples, data);
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;

      break;
      
    case PART_DYNAMIC:

      samples = (float*) malloc( part_size );
      
      //
      // Create the task list
      //
      
      // Initialize these to 0
      start_width = 0;
      number_tasks = 0;
            
      // Setup the task queue
      for (x = block_width; x < width; x = x + block_width) {
        start_height = 0;
        for (y = block_height; y < height; y = y + block_height) {
          if (is_first_task == 1) {
            my_tasks = create_task_list(start_height, start_width, y-1, x-1);
            start_height += block_height;
            is_first_task = 0;
          } else {
            my_tasks = add_task (my_tasks, start_height, start_width, y-1, x-1);
            start_height += block_height;
          }
          number_tasks++;
        }
        // Take care of any extra
        if (start_height < height) {
          my_tasks = add_task (my_tasks, start_height, start_width, height-1, x-1);
          number_tasks++;
        }
        start_width  += block_width;
      }
      
      // Take care of any extra
      if (start_width < width) {
        start_height = 0;      
        for (y = block_height; y < height; y = y + block_height) {
          my_tasks = add_task (my_tasks, start_height, start_width, y-1, width-1);
          start_height += block_height;
          number_tasks++;
        }
        // Take care of any extra
        if (start_height < height) {
          my_tasks = add_task (my_tasks, start_height, start_width, height-1, width-1);
          number_tasks++;
        }
      }

      
      int rank_who_needs_work = 0;
      int* send_arr = (int*) malloc (4 * sizeof(int));
      comm_time = 0;
      for (iter = my_tasks; iter != NULL; iter = iter->next) {
        MPI_Recv(&rank_who_needs_work,1,MPI_INT,MPI_ANY_SOURCE,TAG_DYN_REQ_WORK, MPI_COMM_WORLD,&status);
        get_array(iter, send_arr);
        comm_start = MPI_Wtime();
        MPI_Send(send_arr, 4, MPI_INT, rank_who_needs_work, TAG_DYN_GIVE_WORK, MPI_COMM_WORLD);
        comm_stop = MPI_Wtime();
        comm_time += comm_stop - comm_start;
      }
      
      // Send the packet with all numbers equal to -1, signifying the done condition
      send_arr[0] = -1; send_arr[1] = -1; send_arr[2] = -1; send_arr[3] = -1;
      
      for (node = 1; node < num_procs; node++) {
        MPI_Send(send_arr, 4, MPI_INT, node, TAG_DYN_GIVE_WORK, MPI_COMM_WORLD);
      }
      
      free (send_arr);
            
      break;
    
  } 
  
  printf ("(%d) Done with computations (%f seconds)\n", rank, comp_time);  
    
  //////////////////////////////////////////////////////////////////////////////
  // Receive the data from the slaves
  //////////////////////////////////////////////////////////////////////////////
   switch (part_mode) {
    case PART_STATIC_STRIPS:
    
      start_width = 0;
      end_width   = width - 1;
      start_height = rank * ceil(height/num_procs);
      // Take into account non-nicely-divisable height / num_procs
      end_height   = (rank == (num_procs-1)) ? height - 1 : (rank+1) * ceil(height/num_procs) - 1;   
     
      comm_start = MPI_Wtime();
      for (node = 1; node < num_procs; node++) {
      
        // The receive index
        recv_addr = width*3*node*ceil(height/num_procs);
        
        // Receive data length is the total size divided by the number of processors
        recv_data_len = (node == num_procs-1) ? 3*width*(height-((node)*ceil(height/num_procs))) : 3*width*(end_height-start_height+1);
      
        MPI_Recv (&samples[recv_addr], recv_data_len, MPI_FLOAT, node, 0, MPI_COMM_WORLD, &status);

      }
      comm_stop = MPI_Wtime();
      comm_time = comm_stop - comm_start;
      
      break;
      
    case PART_STATIC_BLOCKS:

      recv_arr     = (float *) malloc (part_size);
      recv_arr_int = (int*) malloc (4*sizeof(int));
      
      comm_start = MPI_Wtime();      
      for (node = 1; node < num_procs; node++) {
        count = get_size (node, data, &my_tasks);
                
        MPI_Recv (recv_arr, count, MPI_FLOAT, node, 0, MPI_COMM_WORLD, &status);
      
        get_array(my_tasks, recv_arr_int);
        start_height = recv_arr_int[0];
        start_width  = recv_arr_int[1];
        end_height   = recv_arr_int[2];
        end_width    = recv_arr_int[3];
        
        count = 0;
        for (x = start_width; x <= end_width; x++) {
          for (y = start_height; y <= end_height; y++) {
            // Check if y is outside range of height
            if (y > height) {
              break;
            }
            for (pixel = 0; pixel < 3; pixel++) {
              // Add to array
              samples[(y * width + x) * 3 + pixel] = recv_arr[count++];
            }
          }
        }
      }
      comm_stop = MPI_Wtime();
      comm_time = comm_stop - comm_start;

      
      free (recv_arr_int);
      free (recv_arr);
    
      break;

    case PART_STATIC_CYCLES:

      recv_arr = (float *) malloc (part_size);
      
      start_width = 0;
      end_width   = width - 1;
      
      comm_start = MPI_Wtime();
      for (node = 1; node < num_procs; node++) {
        MPI_Recv (recv_arr, part_size, MPI_FLOAT, node, 0, MPI_COMM_WORLD, &status);
        // Put the data in recv_arr back in samples

        // Need to put the samples back into samples
        count = 0;
        for (x = start_width; x <= end_width; x++) {
          start_height = node * strip_height;
          while (start_height < height) {
            for (y = start_height; y < start_height + strip_height; y++) {
              if (y > height) {
                break;
              }
              for (pixel = 0; pixel < 3; pixel++) {
                // Add to array
                samples[(y * width + x) * 3 + pixel] = recv_arr[count++];
              }
            }
            start_height = start_height + num_procs*strip_height;
          }
        }

      }
      comm_stop = MPI_Wtime();
      comm_time = comm_stop - comm_start;

      free (recv_arr);

      break;
      
    case PART_DYNAMIC:
    
      recv_arr_int = (int*) malloc (4*sizeof(int));
      recv_arr     = (float *) malloc (part_size);
   
      int num_task_count = 0;
	  
      comm_start = MPI_Wtime();
  
      for (iter = my_tasks; iter != NULL; iter = iter->next) {
      
        if (num_task_count == number_tasks) {
          break;
        }
      
        MPI_Recv (recv_arr_int, 4, MPI_INT, MPI_ANY_SOURCE, TAG_DYNAMIC_DATA_SEND1, MPI_COMM_WORLD, &status);
                
        start_height = recv_arr_int[0];
        start_width  = recv_arr_int[1];
        end_height   = recv_arr_int[2];
        end_width    = recv_arr_int[3];  
        
        MPI_Recv (recv_arr, 3*(end_height-start_height+1)*(end_width-start_width+1), MPI_FLOAT, status.MPI_SOURCE, TAG_DYNAMIC_DATA_SEND2, MPI_COMM_WORLD, &status);
                
        count = 0;
        for (x = start_width; x <= end_width; x++) {
          for (y = start_height; y <= end_height; y++) {
            for (pixel = 0; pixel < 3; pixel++) {
              // Add to array
              samples[(y * width + x) * 3 + pixel] = recv_arr[count++];
            }
          }
        }        
        num_task_count++;
      }
      comm_stop = MPI_Wtime();
      comm_time += comm_stop - comm_start;
      
      free (recv_arr);
      free (recv_arr_int);
      
      printf ("(%d) Data has been received", rank);
      
      break;
    
  }
  
  // Get the longest running node, use that value for the comp_time
  printf ("\n===========================================================\n");
  
  total_comp_time = 0;
  
  for (node = 1; node < num_procs; node++) {
    MPI_Recv (&tmp_comp_time, 1, MPI_DOUBLE, node, 0, MPI_COMM_WORLD, &status);
    total_comp_time += tmp_comp_time;
    if (tmp_comp_time > comp_time) {
      printf ("The longest running node has been swapped from %d (%f) to %d (%f)\n",
               longest_running_node, comp_time, node, tmp_comp_time);
      comp_time = tmp_comp_time;
      longest_running_node = node;
    }
  }
  printf ("===========================================================\n\n");

  //////////////////////////////////////////////////////////////////////////////
  // Save the rendered image.
  //////////////////////////////////////////////////////////////////////////////
  render_name = getRenderName();

  if (!savePixels(render_name, samples, data->width, data->height)) {
    fprintf(stderr, "Failed to save pixel values!\n");
    MPI_Abort( MPI_COMM_WORLD, MPI_ERR_UNKNOWN );
  }

  free(samples);

  printf ("Computation Time:   %.2f seconds from node %d.\n", comp_time, longest_running_node);
  printf ("Total Comp Time:    %.2f seconds.\n", total_comp_time);
  printf ("Communication Time: %.2f seconds.\n", comm_time);
  printf ("C-to-C Ratio:       %.2f\n", (float)comm_time / total_comp_time);

  printf ("Number of processors: %d, PartMode: %d, Block Width: %d, Block Height: %d\nStrip Height: %d, Height: %d, Width: %d, Render Name: %s\n",
          num_procs, part_mode, block_width, block_height,
          strip_height, height, width, render_name);
           
  free(render_name);

  // Everything went fine.
  return;
}
